"""
Various methods used in tests.py, separated to make tests.py neat.

See tests.py for more details.

Most of codes here are essentially part of those in the paper 

    A Census of Exceptional Fillings
        by N. Dunfield, Characters in low-dimensional topology, Contemp. Math. 760 (2020) 143-155,

with slight changes such as minor bug fixes and adaption to up-to-date SnapPy and Regina.
"""

from snappy_10_tets import snappy
# version of snappy used: 3.2

import regina
# version of regina used: 7.3

import re
import networkx as nx

def closed_isosigs(snappy_manifold, trys = 200, max_tets=50, verbose = False):
    """
    Generate a slew of 1-vertex triangulations of a closed manifold
    using SnapPy.
    
    >>> M = snappy.Manifold('m004(1,2)')
    >>> len(closed_isosigs(M, trys=5)) > 0
    True
    """

    M = snappy.Manifold(snappy_manifold)
    assert M.cusp_info('complete?') == [False]
    surgery_descriptions = [M.copy()]

    try:
        for curve in M.dual_curves():
            N = M.drill(curve)
            N.dehn_fill((1,0), 1)
            surgery_descriptions.append(N.filled_triangulation([0]))
    except snappy.SnapPeaFatalError:
        pass

    if len(surgery_descriptions) == 1:
        # Try again, but unfill the cusp first to try to find more
        # dual curves.
        try:
            filling = M.cusp_info(0).filling
            N = M.copy()
            N.dehn_fill((0, 0), 0)
            N.randomize()
            for curve in N.dual_curves():
                D = N.drill(curve)
                D.dehn_fill([filling, (1,0)])
                surgery_descriptions.append(D.filled_triangulation([0]))
        except snappy.SnapPeaFatalError:
            pass

    ans = set()
    for N in surgery_descriptions:
        for i in range(trys):
            T = N.filled_triangulation()
            if T._num_fake_cusps() == 1:
                n = T.num_tetrahedra()
                if n <= max_tets:
                    ans.add((n, T.triangulation_isosig(decorated=False)))
            N.randomize()
            if verbose:
                print('\033[K' + f'closed_isosigs, Trial: {i+1}, Ans number: {len(ans)}', end = '\r')

    if verbose:
        print()
    return [iso for n, iso in sorted(ans)]


def appears_hyperbolic(M):
    acceptable = ['all tetrahedra positively oriented',
                  'contains negatively oriented tetrahedra']
    return M.solution_type() in acceptable and M.volume() > 0

def children(packet):
    child = packet.firstChild()
    while child:
        yield child
        child = child.nextSibling()

def to_regina(data):
    if hasattr(data, '_to_string'):
        data = data._to_string()
    if isinstance(data, str):
        if data.find('(') > -1:
            data = closed_isosigs(data)[0]
        return regina.Triangulation3(data)
    assert isinstance(data, regina.Triangulation3)
    return data

def extract_vector(surface):
    """
    Extract the raw vector of the (almost) normal surface in Regina's
    NS_STANDARD coordinate system.
    """
    S = surface
    T = S.triangulation()
    n = T.countTetrahedra()
    ans = []
    for i in range(n):
        for j in range(4):
            ans.append(S.triangles(i, j))
        for j in range(3):
            ans.append(S.quads(i, j))
    A = regina.NormalSurface(T, regina.NS_STANDARD, ans)
    assert A.sameSurface(S)
    return ans

def haken_sum(S1, S2):
    T = S1.triangulation()
    assert S1.locallyCompatible(S2)
    v1, v2 = extract_vector(S1), extract_vector(S2)
    sum_vec = [x1 + x2 for x1, x2 in zip(v1, v2)]
    A = regina.NormalSurface(T, regina.NS_STANDARD, sum_vec)
    assert S1.locallyCompatible(A) and S2.locallyCompatible(A)
    assert S1.eulerChar() + S2.eulerChar() == A.eulerChar()
    return A


def census_lookup(regina_tri):
    """
    Should the input triangulation be in Regina's census, return the
    name of the manifold, dropping the triangulation number.
    """
    hits = regina.Census.lookup(regina_tri)
    hit = hits[0]
    if hit is not None:
        name = hit.name()
        match = re.search('(.*) : #\d+$', name)
        if match:
            return match.group(1)
        else:
            return match

def standard_lookup(regina_tri):
    match = regina.StandardTriangulation.recognise(regina_tri)
    if match:
        return match.manifold()

def best_match(matches):
    """
    Prioritize the most concise description that Regina provides to
    try to avoid things like the Seifert fibered space of a node being
    a solid torus or having several nodes that can be condensed into a
    single Seifert fibered piece.
    """
    
    def score(m):
        if isinstance(m, regina.SFSpace):
            s = 0
        elif isinstance(m, regina.GraphLoop):
            s = 1
        elif isinstance(m, regina.GraphPair):
            s = 2
        elif isinstance(m, regina.GraphTriple):
            s = 3
        elif m is None:
            s = 10000
        else:
            s = 4
        return (s, str(m))
    return min(matches, key=score)

def identify_with_torus_boundary(regina_tri):
    """
    Use the combined power of Regina and SnapPy to try to give a name
    to the input manifold.
    """
    
    kind, name = None, None
    
    P = regina_tri.clone()
    P.finiteToIdeal()
    P.intelligentSimplify()
    M = snappy.Manifold(P.isoSig())
    M.simplify()
    if appears_hyperbolic(M):
        for i in range(100):
            if M.solution_type() == 'all tetrahedra positively oriented':
                break
            M.randomize()
        
        if not M.verify_hyperbolicity(bits_prec=100):
            raise RuntimeError('Cannot prove hyperbolicity for ' +
                               M.triangulation_isosig())
        kind = 'hyperbolic'
        ids = M.identify()
        if ids:
            name = ids[0].name()
    else:
        match = standard_lookup(regina_tri)
        if match is None:
            Q = P.clone()
            Q.idealToFinite()
            Q.intelligentSimplify()
            match = standard_lookup(Q)
        if match is not None:
            kind = match.__class__.__name__
            name = str(match)
        else:
            name = P.isoSig()
    return kind, name

def is_toroidal(regina_tri):
    """
    Checks for essential tori and returns the pieces of the
    associated partial JSJ decomposition.
    
    >>> T = to_regina('hLALAkbccfefgglpkusufk')  # m004(4,1)
    >>> is_toroidal(T)[0]
    True
    >>> T = to_regina('hvLAQkcdfegfggjwajpmpw')  # m004(0,1)
    >>> is_toroidal(T)[0]
    True
    >>> T = to_regina('nLLLLMLPQkcdgfihjlmmlkmlhshnrvaqtpsfnf')  # 5_2(10,1)
    >>> T.isHaken()
    True
    >>> is_toroidal(T)[0]
    False

    Note: currently checks all fundamental normal tori; possibly
    the theory lets one just check *vertex* normal tori.
    """
    T = regina_tri
    assert T.isZeroEfficient()
    surfaces = regina.NNormalSurfaceList.enumerate(T,
                          regina.NS_QUAD, regina.NS_FUNDAMENTAL)
    for i in range(surfaces.size()):
        S = surfaces.surface(i)
        if S.eulerChar() == 0:
            if not S.isOrientable():
                S = S.doubleSurface()
            assert S.isOrientable()
            X = S.cutAlong()
            X.intelligentSimplify()
            X.splitIntoComponents()
            pieces = list(children(X))
            if all(not C.hasCompressingDisc() for C in pieces):
                ids = [identify_with_torus_boundary(C) for C in pieces]
                return (True, sorted(ids))
                
    return (False, None)


def decompose_along_tori(regina_tri):
    """
    First, finds all essential normal tori in the manifold associated
    with fundamental normal surfaces.  Then takes a maximal disjoint
    collection of these tori, namely the one with the fewest tori
    involved, and cuts the manifold open along it.  It tries to
    identify the pieces, removing any (torus x I) components. 

    Returns: (has essential torus, list of pieces)

    Note: This may fail to be the true JSJ decomposition because there
    could be (torus x I)'s in the list of pieces and it might well be
    possible to amalgamate some of the pieces into a single SFS.
    """
    
    T = regina_tri
    assert T.isZeroEfficient()
    essential_tori = []
    surfaces = regina.NNormalSurfaceList.enumerate(T,
                          regina.NS_QUAD, regina.NS_FUNDAMENTAL)
    for i in range(surfaces.size()):
        S = surfaces.surface(i)
        if S.eulerChar() == 0:
            if not S.isOrientable():
                S = S.doubleSurface()
            assert S.isOrientable()
            X = S.cutAlong()
            X.intelligentSimplify()
            X.splitIntoComponents()
            pieces = list(children(X))
            if all(not C.hasCompressingDisc() for C in pieces):
                essential_tori.append(S)

    if len(essential_tori) == 0:
        return False, None
    
    D = nx.Graph()
    for a, A in enumerate(essential_tori):
        for b, B in enumerate(essential_tori):
            if a < b:
                if A.disjoint(B):
                    D.add_edge(a, b)

    cliques = list(nx.find_cliques(D))
    if len(cliques) == 0:
        clique = [0]
    else:
        clique = min(cliques, key=len)
    clique = [essential_tori[c] for c in clique]
    A = clique[0]
    for B in clique[1:]:
        A = haken_sum(A, B)

    X = A.cutAlong()
    X.intelligentSimplify()
    X.splitIntoComponents()
    ids = [identify_with_torus_boundary(C) for C in list(children(X))]
    # Remove products
    ids = [i for i in ids if i[1] not in ('SFS [A: (1,1)]', 'A x S1')]
    return (True, sorted(ids))

def regina_name(closed_snappy_manifold, trys=100):
    """
    >>> regina_name('m004(1,0)')
    'S3'
    >>> regina_name('s006(-2, 1)')
    'SFS [A: (5,1)] / [ 0,-1 | -1,0 ]'
    >>> regina_name('m010(-1, 1)')
    'L(3,1) # RP3'
    >>> regina_name('m022(-1,1)')
    'SFS [S2: (3,2) (3,2) (4,-3)]'
    >>> regina_name('v0004(0, 1)')
    'SFS [S2: (2,1) (4,1) (15,-13)]'
    >>> regina_name('m305(1, 0)')
    'L(3,1) # RP3'
    """
    M = closed_snappy_manifold
    isosigs = closed_isosigs(M, trys=trys, max_tets=100)
    if len(isosigs) == 0:
        return
    T = to_regina(isosigs[0])
    if T.isIrreducible() or len(T.summands()) == 1: # if T is irreducible, T.isIrreducible will be faster than T.summands (I think)
        if T.countTetrahedra() <= 11:
            for i in range(3):
                T.simplifyExhaustive(i)
                name = census_lookup(T)
                if name is not None:
                    return name
            
        matches = [standard_lookup(to_regina(iso)) for iso in isosigs]
        match = best_match(matches)
        if match is not None:
            return str(match)
    else:
        piece_list = T.summands()
        pieces = [regina_name(P) for P in piece_list]
        if None not in pieces:
            return ' # '.join(sorted(pieces))
        
